import os

import numpy as np
import torch
from PIL import Image
from tqdm import tqdm
from bezier_curve import bezier_curve

def norm11(x):
    x = (x - x.min()) / (x.max() - x.min())
    x = x*2-1
    return x

def nonlinear_transformation(slices):
    assert slices.max() <= 1, print(slices.max())
    assert slices.min() >= -1

    points_1 = [[-1, -1], [-1, -1], [1, 1], [1, 1]]
    xvals_1, yvals_1 = bezier_curve(points_1, nTimes=100000)
    xvals_1 = np.sort(xvals_1)

    points_2 = [[-1, -1], [-0.5, 0.5], [0.5, -0.5], [1, 1]]
    xvals_2, yvals_2 = bezier_curve(points_2, nTimes=100000)
    xvals_2 = np.sort(xvals_2)
    yvals_2 = np.sort(yvals_2)

    points_3 = [[-1, -1], [-0.5, 0.5], [0.5, -0.5], [1, 1]]
    xvals_3, yvals_3 = bezier_curve(points_3, nTimes=100000)
    xvals_3 = np.sort(xvals_3)

    points_4 = [[-1, -1], [-0.75, 0.75], [0.75, -0.75], [1, 1]]
    xvals_4, yvals_4 = bezier_curve(points_4, nTimes=100000)
    xvals_4 = np.sort(xvals_4)
    yvals_4 = np.sort(yvals_4)

    points_5 = [[-1, -1], [-0.75, 0.75], [0.75, -0.75], [1, 1]]
    xvals_5, yvals_5 = bezier_curve(points_5, nTimes=100000)
    xvals_5 = np.sort(xvals_5)

    """
    slices, nonlinear_slices_2, nonlinear_slices_4 are source-similar images
    nonlinear_slices_1, nonlinear_slices_3, nonlinear_slices_5 are source-dissimilar images
    """
    nonlinear_slices_1 = np.interp(slices, xvals_1, yvals_1)
    nonlinear_slices_1[nonlinear_slices_1 == 1] = -1
    
    nonlinear_slices_2 = np.interp(slices, xvals_2, yvals_2)

    nonlinear_slices_3 = np.interp(slices, xvals_3, yvals_3)
    nonlinear_slices_3[nonlinear_slices_3 == 1] = -1

    nonlinear_slices_4 = np.interp(slices, xvals_4, yvals_4)

    nonlinear_slices_5 = np.interp(slices, xvals_5, yvals_5)
    nonlinear_slices_5[nonlinear_slices_5 == 1] = -1

    return slices, nonlinear_slices_1, nonlinear_slices_2, \
           nonlinear_slices_3, nonlinear_slices_4, nonlinear_slices_5


def make_ds(dataset):
    dir=f'data/{dataset}-128-160-my'
    for f in ['ss', 'sd']:
        for n in range(3):
            if not os.path.exists(os.path.join(f'data/{dataset}-128-160-my-ds', f'{f}_{n}')):
                os.makedirs(os.path.join(f'data/{dataset}-128-160-my-ds', f'{f}_{n}'))

    for img_name in tqdm(os.listdir(os.path.join(dir, 'images'))[:]):
        slices = np.array(Image.open(os.path.join(dir, 'images', img_name)).convert('L'))/255

        if slices.max()>0:
            slices = norm11(slices)
        else:
            slices -= 1

        slices, nonlinear_slices_1, nonlinear_slices_2, \
        nonlinear_slices_3, nonlinear_slices_4, nonlinear_slices_5 = nonlinear_transformation(slices)

        for n, x in enumerate([slices, nonlinear_slices_2, nonlinear_slices_4]):
            Image.fromarray(np.uint8((x+1)/2*255)).save(os.path.join(f'data/{dataset}-128-160-my-ds', f'ss_{n}', img_name))

        for n, x in enumerate([nonlinear_slices_1, nonlinear_slices_3, nonlinear_slices_5]):
            Image.fromarray(np.uint8((x+1)/2*255)).save(os.path.join(f'data/{dataset}-128-160-my-ds', f'sd_{n}', img_name))


make_ds('CANDI')
make_ds('OASIS')